/**
* stego.c
* Authors: Darren Castillo, Rudolph Carvell Scott II, John Torres
*/

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <unistd.h>

//constants taken from CS4953 bitmap c++
#define BI_RGB        0L        // (no compression)
#define BI_RLE8       1L        // (run-length encoding, 8 bits)
#define BI_RLE4       2L        // (run-length encoding, 4 bits)
#define BI_BITFIELDS  3L

#define SEEDSIZEINBITS 11
#define MAXSZ (1 << sizeof(int)-SEEDSIZEINBITS)-1

typedef unsigned int TYPE_UINT32;
typedef unsigned short TYPE_UINT16;
typedef int TYPE_INT32;
typedef unsigned char TYPE_UCHAR8;

//###the following structure information is taken from wingdi.h###
//Bitmap File Header
typedef struct tagBITMAPFILEHEADER{ //14 bytes
    TYPE_UINT16    bfType;
    TYPE_UINT32   bfSize;      // total length of bitmap file
    TYPE_UINT32   stegokey;     //first SEEDSIZEINBITS bits: random seed.
                               //other bits: Message size
    TYPE_UINT32   bfOffBits;   // offset to start of actual pixel data
} BITMAPFILEHEADER;
//DIB Header
typedef struct tagBITMAPINFOHEADER{ //40 bytes
    TYPE_UINT32   biSize;              // size of BITMAPINFOHEADER
    TYPE_INT32    biWidth;             // width in pixels
    TYPE_INT32    biHeight;            // height in pixels
    TYPE_UINT16    biPlanes;            // always 1
    TYPE_UINT16    biBitCount;          // color bits per pixel
    TYPE_UINT32   biCompression;       // BI_RGB, BI_RLE8, BI_RLE4
    TYPE_UINT32   biSizeImage;         // total bytes in image
    TYPE_INT32    biXPelsPerMeter;     // 0, or optional horizontal resolution
    TYPE_INT32    biYPelsPerMeter;     // 0, or optional vertical resolution
    TYPE_UINT32   biClrUsed;           // colors actually used (normally zero, can be lower than biBitCount)
    TYPE_UINT32   biClrImportant;      // important colors actualy used (normally zero)
} BITMAPINFOHEADER;
typedef struct tagRGBQUAD{
    TYPE_UCHAR8    rgbBlue;
    TYPE_UCHAR8    rgbGreen;
    TYPE_UCHAR8    rgbRed;
} COLOR;

static void displayFileInfo(char *pFileName, BITMAPFILEHEADER *pFileHdr,
                            BITMAPINFOHEADER *pFileInfo, 
                            COLOR *pColorTable, 
                            unsigned char *pData);
static void readFileData(char *file, unsigned int *fileSize, unsigned char **image);
static unsigned int writeStegoImage(unsigned char *image, unsigned int fileSize, char *fileName);
static unsigned int writeExtracted(unsigned char *message, unsigned int fileSize, char *fileName);
static TYPE_INT32 hidedata(TYPE_UCHAR8 *data, TYPE_UCHAR8 *message,
                           TYPE_UCHAR8 key,
                           TYPE_INT32 messageSizeBytes,
                           unsigned long long numPixels);
static TYPE_UCHAR8 *extractdata(BITMAPFILEHEADER *fileHeader, TYPE_UCHAR8 *data);
static TYPE_UCHAR8 decideLSBCount(TYPE_UCHAR8 value, int numbits);
static TYPE_UCHAR8 getBits(TYPE_UCHAR8 value, int n);
static TYPE_UINT32 getCapacity(TYPE_UCHAR8 *data, TYPE_UCHAR8 key, unsigned long long numPixels);

//needed by multiple functions to account for padding bytes
TYPE_INT32 imgWidth;

int main(int argc, char *argv[]){
   unsigned char *image, *imagePixelData, *message, flag;
   unsigned int hiddenSz, fileSz, messageSz;
   unsigned short seedVal;
   unsigned long long numPix;
   BITMAPFILEHEADER *bmpHeader;
   BITMAPINFOHEADER *bmpInfoHeader;
   COLOR *bmpColorTable;

   //Check Arguments
   if(argc < 3){
      fprintf(stderr,"Usage: hideex <COVER.bmp> <HIDDENDATA.*> <-x | INT seed>\n");
      fprintf(stderr,"The flag -x will make the program attempt to extract a hidden message from the file COVER.\n");
      fprintf(stderr,"If this flag is not present, the program attempts to hide whatever data is given by HIDDENDATA within COVER.\n");
      exit(1);
   }
   
   //Check for extraction flag ("-x") or for random seed (int)
   if(strcmp(argv[3], "-x") == 0){
      flag = 1; //flag indicates if extracting or embedding
   }else{
      seedVal = atoi(argv[3]);
      seedVal &= (1 << SEEDSIZEINBITS) - 1; //Force the random seed to fit in 11 bits.
      flag = 0;
   }

   //read in cover image, store to image
   readFileData(argv[1], &fileSz, &image);
   
   //begin reading in bitmap headers here
   bmpHeader = (BITMAPFILEHEADER *) (image - 2);
   bmpInfoHeader = (BITMAPINFOHEADER *) (image + 14);
   bmpColorTable = (COLOR *) (image + 14 + bmpInfoHeader->biSize);
   imagePixelData = image + bmpHeader->bfOffBits;

   imgWidth = bmpInfoHeader->biWidth;
   
   //This ensures that the file passed in is a bitmap.
   char *indexOfFileType = strstr((char *) image, "BM");
   if(indexOfFileType != (char *) image){
      fprintf(stderr, "%s is not a Windows Bitmap.\n", argv[1]);
      exit(1);
   }else if(bmpInfoHeader->biBitCount != 24){
      fprintf(stderr, "%s is not 24-bit bmp\n", argv[1]);
      exit(1);
   }
   else
      displayFileInfo(argv[1], bmpHeader, bmpInfoHeader, bmpColorTable, imagePixelData);

   if(flag){//extracting
      //extract message from stegoimage
      message = extractdata(bmpHeader, imagePixelData);
      printf("Writing to file \"%s\"\n", argv[2]);
      writeExtracted(message, (bmpHeader->stegokey >> SEEDSIZEINBITS), argv[2]);
      free(message);
   }
   else{//embedding
      //read message, store to message
      readFileData(argv[2], &messageSz, &message);
      numPix = bmpInfoHeader->biWidth * bmpInfoHeader->biHeight;
      
      //embed message into cover
      hiddenSz = hidedata(imagePixelData, message, seedVal, messageSz, numPix);
      
      //Show capacity info
      TYPE_UINT32 maxCapacity = getCapacity(imagePixelData,seedVal,numPix);
      printf("Capacity in bytes: %d\n", maxCapacity/8);
      printf("Bytes Hidden: %u\n",hiddenSz);
      printf("Usage: %.2f%%\n", hiddenSz*100.0/(maxCapacity/8.0));

      //insert the message size and random seed into the 
      bmpHeader->stegokey = hiddenSz;
      bmpHeader->stegokey = bmpHeader->stegokey << SEEDSIZEINBITS;
      bmpHeader->stegokey = bmpHeader->stegokey | seedVal;

      //write out data
      writeStegoImage(image, fileSz, argv[1]);
   }

   return 0;
}
/**
 * Prints generic statistical data about the supplied cover.
 * */
static void displayFileInfo(char *pFileName, BITMAPFILEHEADER *pFileHdr,
            BITMAPINFOHEADER *pFileInfo, COLOR *pColorTable, unsigned char *pData)
{
   int numColors;
   
   printf("\nFile Info for %s: \n\n", pFileName);
   printf("File Type: BM\n");
   printf("File Size: %d\nData Offset: %x\n\n", pFileHdr->bfSize, (pFileHdr->bfOffBits));

   numColors = 1 << (pFileInfo->biBitCount);
   
   printf("Bit Map Image Info: \n\nSize: %d\nWidth: %d\nHeight: %d\nPlanes: %d\n"
        "Bits/Pixel: %d (%d colors)\n"
        "Compression: %d\nImage Size: %d\nRes X: %d\nRes Y: %d\nColors: %d\nImportant Colors: %d\n\n",
        pFileInfo->biSize,
        pFileInfo->biWidth,
        pFileInfo->biHeight,
        pFileInfo->biPlanes,
        pFileInfo->biBitCount, numColors,
        pFileInfo->biCompression,
        pFileInfo->biSizeImage,
        pFileInfo->biXPelsPerMeter,
        pFileInfo->biYPelsPerMeter,
        pFileInfo->biClrUsed,
        pFileInfo->biClrImportant);
}
/**
 * Read in the cover data from <file>, store the data to <image>, store the file size in bytes. <obvious>
 * */
static void readFileData(char *file, unsigned int *fileSize, unsigned char **image){
   FILE *infile;

   if((infile = fopen(file,"r")) == NULL){
      fprintf(stderr,"error opening file %s\n",file);
      perror(NULL);
      exit(1);
   }

   //A little trick to obtain the filesize
   fseek(infile,0,SEEK_END);
   *fileSize = ftell(infile);
   fseek(infile,0,SEEK_SET);

   if((*image = (unsigned char *) malloc(*fileSize)) == NULL){
      perror("malloc error");
      exit(1);
   }
   fread(*image, sizeof(unsigned char), *fileSize, infile);
   fclose(infile);
}
/**
 * Writes a stego-image of size <fileSize>(bytes). The original cover's is passed in as <fileName>. <image> contains stego-data. 
 * */
static unsigned int writeStegoImage(unsigned char *image, unsigned int fileSize, char *fileName){
   FILE *outfile;
   char fileBuffer[256];
   int offset;
   unsigned int bytesWritten;

   //add .hx.bmp to stego cover
   offset = strlen(fileName)-4;
   strncpy(fileBuffer, fileName, offset);
   fileBuffer[offset] = '\0';
   strcat(fileBuffer, ".hx.bmp");

   //begin write
   if((outfile = fopen(fileBuffer, "w+"))==NULL){
      fprintf(stderr,"error writing to %s\n",fileBuffer);
      perror(NULL);
      exit(1);
   }
   if((bytesWritten = fwrite(image, sizeof(unsigned char), fileSize, outfile)) != fileSize){
      fprintf(stderr, "BMP: incorrect write to file. bytes to write: %d, bytes written: %d\n", fileSize, bytesWritten);
      perror("incorrect write to file\n");
      exit(1);
   }
   fclose(outfile);

   return bytesWritten;
}
/**
 * Writes extracted data <message> to <fileName. <fileSize> is needed to indicate how much to write in bytes.
 * */
static unsigned int writeExtracted(unsigned char *message, unsigned int fileSize, char *fileName){
    FILE *outfile;
    char fileBuffer[256];
    TYPE_UCHAR8* byteToWrite;
    unsigned int bytesWritten;
    unsigned long long int totalWritten;

    //begin write
    if((outfile = fopen(fileName, "w+"))==NULL){
        fprintf(stderr,"error writing to %s\n",fileBuffer);
        perror(NULL);
        exit(1);
    }
    byteToWrite = message;
    totalWritten = 0;
    while(totalWritten < fileSize) {
        if((bytesWritten = fwrite(byteToWrite, sizeof(unsigned char), 1, outfile)) != 1){
            fprintf(stderr, "Extract: incorrect write to file. bytes to write: %d, bytes written: %d\n", fileSize, bytesWritten);
            exit(1);
        }
        totalWritten += bytesWritten;
        byteToWrite++;
    }
    fclose(outfile);
    return bytesWritten;
}
/**
 * Creates the stegomessage, using <data> as cover, and <message> as the message. messageSizeBytes is in bytes.
 * */
static TYPE_INT32 hidedata(TYPE_UCHAR8 *data, TYPE_UCHAR8 *message, TYPE_UCHAR8 key, TYPE_INT32 messageSizeBytes,
                     unsigned long long numPixels){
   int numbits, numLSB;
   TYPE_INT32 msgIndex, totalHidden;
   unsigned long long pxIndex;
   TYPE_UINT16 bitbuf;
   TYPE_UINT32 pixelCounter;
   COLOR *pixelBuf =(COLOR *) data; // represents current pixel
   TYPE_UCHAR8 *padOffset; // Used to offset COLOR pixel if hits padding
   TYPE_UCHAR8 hidebuf; // represents current message
   TYPE_UCHAR8 bitOffset; // used to replace n bits to 0
   TYPE_UCHAR8 numPadding; // number of padding bits in a row
   TYPE_UCHAR8 r, g, b; // Holds RGB value to use in condtions and calculations
   totalHidden = 0; // number of bits hidden
   msgIndex = 0; // current index of info
   pxIndex = 0; // current index of pixel
   bitbuf = 0; // bits of data that are ready to be hidden
   numbits = 0; // number of bits in bitbuf
   pixelCounter = 1; // Pixels read in per row

   srand(key); // seed rand to key

   do {
      // Read in Hidden Data -- Combine next byte of info into bitbuf
      // Note: One byte can only be read in, and 3~5 bits may be hidden
      if(numbits < 5 && msgIndex < messageSizeBytes) {
         hidebuf = (TYPE_UCHAR8) *(message + msgIndex); // gets next byte from the msg
         bitbuf = (bitbuf << 8) + hidebuf; // combines bits into a single stream
         numbits += 8;
         msgIndex++;
      }

      // assign to shorten code
      r = pixelBuf->rgbRed;
      g = pixelBuf->rgbGreen;
      b = pixelBuf->rgbBlue;

      // Choose indicator channel (case) -- 0,1,2 is R,G,B respectively
      // Replace LSB: leftmost channel is 0, and rightmost is 1
      switch(rand() % 3) {
         case 0: // RED
            if(g <= b) { // decide which of the 2 remaining channels is lower
               // write RED + indicator
               pixelBuf->rgbRed = r & 254; // LSB=0 (left channel is indicator)
               // write GREEN + hidden
               numLSB = decideLSBCount(g, numbits);
               bitOffset = 256 - (1 << numLSB);// (1 << numLSB); // Ex. n=3 11111111-0000111
               pixelBuf->rgbGreen = (g&bitOffset)+(bitbuf>>(numbits-numLSB));
            }
            else {
               // write RED + indicator
               pixelBuf->rgbRed = r | 1; // LSB=1 (right channel is indicator)
               // write BLUE + hidden
               numLSB = decideLSBCount(b, numbits);
               bitOffset = 256 - (1 << numLSB);
               pixelBuf->rgbBlue = (b&bitOffset)+(bitbuf>>(numbits-numLSB));
            }
            break;
         case 1: // GREEN
            if(r < b) {
               // write RED + hidden
               numLSB = decideLSBCount(r, numbits);
               bitOffset = 256 - (1 << numLSB);
               pixelBuf->rgbRed = (r&bitOffset)+(bitbuf>>(numbits-numLSB));
               // write GREEN + indicator
               pixelBuf->rgbGreen = g & 254; // LSB is 0
            }
            else {
               // write GREEN + indicator
               pixelBuf->rgbGreen = g | 1; // LSB is 1
               // write BLUE + hidden
               numLSB = decideLSBCount(b, numbits);
               bitOffset = 256 - (1 << numLSB);
               pixelBuf->rgbBlue = (b&bitOffset)+(bitbuf>>(numbits-numLSB));
            }
            break;
         case 2: // BLUE
            if(r <= g) {
               // write RED + hidden
               numLSB = decideLSBCount(r, numbits);
               bitOffset = 256 - (1 << numLSB);
               pixelBuf->rgbRed = (r&bitOffset)+(bitbuf>>(numbits-numLSB));
               // write BLUE + indicator
               pixelBuf->rgbBlue = b & 254; // LSB is 0
            }
            else {
               // write GREEN + hidden
               numLSB = decideLSBCount(g, numbits);
               bitOffset = 256 - (1 << numLSB);
               pixelBuf->rgbGreen = (g&bitOffset)+(bitbuf>>(numbits-numLSB));
               // write BLUE + indicator
               pixelBuf->rgbBlue = b | 1; // LSB is 1
            }
      }
      // Update bitbuf
      bitbuf = bitbuf << (16-numbits+numLSB); // Separated. If on a single line
      bitbuf = bitbuf >> (16-numbits+numLSB); //    even with ()'s, cancels out
      numbits -= numLSB;
      totalHidden += numLSB;

      // End of row padding offset
      if((imgWidth*3)%4 != 0 && (pxIndex+1) % imgWidth == 0) {
         numPadding = 4-((imgWidth)*3)%4;
         padOffset = (TYPE_UCHAR8*)pixelBuf; // To offset a COLOR by number of padding bytes
         padOffset += numPadding;
         pixelBuf = (COLOR*)padOffset; // Convert back;
         pixelCounter = 0; // Reset counter for next row
      }

      pxIndex++;
      pixelBuf++;
      pixelCounter++;
   }
   while(pxIndex < numPixels && ((msgIndex) < messageSizeBytes || numbits > 0));

   if((msgIndex+1) < messageSizeBytes)
       printf("NOTE: Not enough capacity to hold all of the hidden message!\n");
   return msgIndex; // Returns number of BYTES hidden
}
/**
 * Extract data from a stego image made with this program. Parameters from main().
 * */
static TYPE_UCHAR8 *extractdata(BITMAPFILEHEADER *fileHeader, TYPE_UCHAR8 *data) {
   int numbits, msgIndex, numLSB, pxIndex;
   TYPE_UINT32 messageSz, pixelCounter, messageSzInBits;
   TYPE_UINT16 bitbuf, numExtracted;
   COLOR *pixelBuf = (COLOR*) data; // represents current pixel
   TYPE_UCHAR8 *padOffset; // Used to offset COLOR pixel if hits padding
   TYPE_UCHAR8 *extracted; // represents current extracted message
   TYPE_UCHAR8 r, g, b; // Holds RGB value to use in condtions and calculations
   TYPE_UCHAR8 numPadding; // number of padding bits in a row
   TYPE_UCHAR8 key;
   numExtracted = 0; // total number of hidden bits extracted
   msgIndex = 0; // current index of info
   pxIndex = 0; // current index of pixel
   bitbuf = 0; // bits of data that are ready to be hidden
   numbits = 0; // number of bits in bitbuf
   pixelCounter = 1; // Pixels read in per row

   // extract paramaters from header
   key = (TYPE_UCHAR8)(fileHeader->stegokey % (TYPE_UINT32) (1 << SEEDSIZEINBITS));
   messageSz = (TYPE_UINT32)fileHeader->stegokey >> SEEDSIZEINBITS;

   // Set Seed
   messageSzInBits = messageSz * 8;
   extracted = malloc(sizeof(TYPE_UCHAR8)*messageSz);
   srand(key);
   do {
      // Extract Hidden Data -- Combine next byte of info into bitbuf
      // Note: One byte can only be read in, and 3~5 bits may be hidden
      if(numbits > 7) {
         extracted[msgIndex] = bitbuf>>(numbits-8); // Stores byte into memory
         bitbuf = bitbuf << (16-numbits+8); // Separated -- if on a single line
         bitbuf = bitbuf >> (16-numbits+8); //    even with ()'s, cancels out
         numbits -= 8;
         msgIndex++;
      }

      // assign to shorten code
      r = pixelBuf->rgbRed;
      g = pixelBuf->rgbGreen;
      b = pixelBuf->rgbBlue;

      // Find indicator channel (case) -- 0,1,2 is R,G,B respectively
      // Replace LSB: leftmost channel is 0, and rightmost is 1
      switch(rand() % 3) {
         case 0: // RED
            if(r % 2 == 0) { // LSB is 0 (Green is hidden channel)
               numLSB = decideLSBCount(g, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(g, numLSB);
            }
            else { // LSB is 1 (Blue is hidden channel)
               numLSB = decideLSBCount(b, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(b, numLSB);
            }
            break;
         case 1: // GREEN
            if(g % 2 == 0) { // LSB is 0 (Red is hidden channel)
               numLSB = decideLSBCount(r, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(r, numLSB);
            }
            else { // LSB is 1 (Blue is hidden channel)
               numLSB = decideLSBCount(b, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(b, numLSB);
            }
            break;
         case 2: // BLUE
            if(b % 2 == 0) { // LSB is 0 (Red is hidden channel)
               numLSB = decideLSBCount(r, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(r, numLSB);
            }
            else { // LSB is 1 (Green is hidden channel)
               numLSB = decideLSBCount(g, messageSzInBits-numExtracted);
               bitbuf = (bitbuf << numLSB) + getBits(g, numLSB);
            }
      } // end switch
      numbits += numLSB;
      numExtracted += numLSB;

      if((imgWidth*3)%4 != 0 && (pxIndex+1) % imgWidth == 0) {
         numPadding = 4-((imgWidth)*3)%4;
         padOffset = (TYPE_UCHAR8*)pixelBuf; // To offset a COLOR by no. of padding bytes
         padOffset += numPadding;
         pixelBuf = (COLOR*)padOffset; // Convert back;
         pixelCounter = 0; // Reset counter for next row
      }

      pxIndex++;
      pixelBuf++;
      pixelCounter++;

      if(numExtracted >= messageSzInBits && numbits > 1)
         extracted[msgIndex] = bitbuf << (8-numbits);

   }
   while(numExtracted < messageSzInBits && msgIndex < messageSz);
   return extracted;
}
/**
 * Determines how many LSBs to replace based on the color channel's <value>.
 * */
static TYPE_UCHAR8 decideLSBCount(TYPE_UCHAR8 value, int numbits){
   // Note: Thresholds were decided through visual perception trial & error
   if(numbits < 4)
      return numbits;
   if(value < 32)
      return 5;
   if(value < 64)
      return 4;
   return 3;
}
/**
 * Used by extraction method to obtain full bytes from partial bytes
 * */
static TYPE_UCHAR8 getBits(TYPE_UCHAR8 value, int numLSB) {
   TYPE_UCHAR8 offset = value << (8-numLSB);
   return offset >> (8-numLSB);
}
/**
 * Get the max capacity of the cover image in bits (since the algorithm can produce partial bytes).
 * */
static TYPE_UINT32 getCapacity(TYPE_UCHAR8 *data, TYPE_UCHAR8 key, unsigned long long numPixels) {
   unsigned long long capacity, pxIndex;
   TYPE_UINT32 pixelCounter;
   COLOR *pixelBuf = (COLOR*) data; // represents current pixel's color values
   TYPE_UCHAR8 *padOffset; // Used to offset COLOR pixel if hit with a end-of-row padding
   TYPE_UCHAR8 r, g, b; // Holds RGB value to use in condtions and calculations
   TYPE_UCHAR8 numPadding; // number of padding bits in a row
   capacity = 0; // the max capactiy of the image
   pxIndex = 0; // index of the pixel
   pixelCounter = 1; // Pixels read in per row
   srand(key); // Set and Hide Key
   
   while(pxIndex < numPixels) {
      // assign to shorten code
      r = pixelBuf->rgbRed;
      g = pixelBuf->rgbGreen;
      b = pixelBuf->rgbBlue;

      // Choose indicator channel (case) -- 0,1,2 is R,G,B respectively
      switch(rand() % 3) {
         case 0: // RED
            if(g <= b) // decide which of the 2 remaining channels is lower
               capacity += decideLSBCount(g, 4); // 4 is used so 1st cond. is not used
            else
               capacity += decideLSBCount(b, 4);
            break;
         case 1: // GREEN
            if(r < b)
               capacity += decideLSBCount(r, 4);
            else
               capacity += decideLSBCount(b, 4);
            break;
         case 2: // BLUE
            if(r <= g)
               capacity += decideLSBCount(r, 4);
            else
               capacity += decideLSBCount(g, 4);
      }
      
      // End of row padding offset
      if((imgWidth*3)%4 != 0 && (pxIndex+1) % imgWidth == 0) {
         numPadding = 4-((imgWidth)*3)%4;
         padOffset = (TYPE_UCHAR8*)pixelBuf; // To offset a COLOR by no. of padding
         padOffset += numPadding;
         pixelBuf = (COLOR*)padOffset; // Convert back;
         pixelCounter = 0; // Reset counter for next row
      }

      pxIndex++;
      pixelBuf++;
      pixelCounter++;
   }
   return capacity;
}
